'''------------------------------------------------------------------------------------------------
 $ @Author: yangshangqi
 $ @Date: 2021-12-20 18:37:03
 $ @LastEditTime: 2021-12-22 20:29:23
 $ @LastEditors: Please set LastEditors
 $ @Description: 工具代码: *计算IOU的函数->get_iou
 $ @FilePath: \YSQTrackerEvaluationToolkit\TOOLS\GetIou.py
--------------------------------------------------------------------------------------------------'''
import numpy as np

'''-------------------------------------------------------------------------------------------
 * @description: 计算两个矩形框的重叠度（IOU）;#!要求输入box必须为一个整型元组-tuple
 * @param {*} bbox1-> 整型元组[x1,y1,x2,y2]:左上角与右下角坐标
 * @param {*} bbox2-> 整型元组[x1,y1,x2,y2]:左上角与右下角坐标
 * @return {*}
---------------------------------------------------------------------------------------------'''
def get_iou(bbox1, bbox2):
    """
    Calculates the intersection-over-union of two bounding boxes.
    """
    bbox1 = [float(x) for x in bbox1]
    bbox2 = [float(x) for x in bbox2]
    (x0_1, y0_1, x1_1, y1_1) = bbox1
    (x0_2, y0_2, x1_2, y1_2) = bbox2
    # get the overlap rectangle
    overlap_x0 = max(x0_1, x0_2)
    overlap_y0 = max(y0_1, y0_2)
    overlap_x1 = min(x1_1, x1_2)
    overlap_y1 = min(y1_1, y1_2)
    # check if there is an overlap
    if overlap_x1 - overlap_x0 <= 0 or overlap_y1 - overlap_y0 <= 0:
        return 0
    # if yes, calculate the ratio of the overlap to each ROI size and the unified size
    size_1 = (x1_1 - x0_1) * (y1_1 - y0_1)
    size_2 = (x1_2 - x0_2) * (y1_2 - y0_2)
    size_intersection = (overlap_x1 - overlap_x0) * (overlap_y1 - overlap_y0)
    size_union = size_1 + size_2 - size_intersection
    return size_intersection / size_union

def getIouArray(bbox1,bbox2):
    x0_1 = bbox1[:,0]
    y0_1 = bbox1[:,1]
    x1_1 = bbox1[:,2]
    y1_1 = bbox1[:,3]

    x0_2 = bbox2[:,0]
    y0_2 = bbox2[:,1]
    x1_2 = bbox2[:,2]
    y1_2 = bbox2[:,3]

    overlap_x0 = np.zeros(x0_1.shape[0])
    overlap_y0 = np.zeros(x0_1.shape[0])
    overlap_x1 = np.zeros(x0_1.shape[0])
    overlap_y1 = np.zeros(x0_1.shape[0])

    flage = np.zeros(x0_1.shape[0])

    for i in range(x0_1.shape[0]):
        overlap_x0[i] = _findMax(x0_1[i], x0_2[i])
        overlap_y0[i] = _findMax(y0_1[i], y0_2[i])
        overlap_x1[i] = _findMin(x1_1[i], x1_2[i])
        overlap_y1[i] = _findMin(y1_1[i], y1_2[i])
        if overlap_x1[i] - overlap_x0[i] > 0 or overlap_y1[i] - overlap_y0[i] > 0:
            flage[i] = 1

    size_1 = (x1_1 - x0_1) * (y1_1 - y0_1)
    size_2 = (x1_2 - x0_2) * (y1_2 - y0_2)
    size_intersection = (overlap_x1 - overlap_x0) * (overlap_y1 - overlap_y0)
    size_union = size_1 + size_2 - size_intersection +1e-9
    ans = (size_intersection / size_union)*flage
    return abs(ans)

def getVectorM(v1,v2,func):
    if(v1.shape[0] != v2.shape[0]):
        print("ERROR:./Tools/GetIou.py/getVectorM::输入向量维度不一致")
        return
    else:
        ans = np.zeros(v1.shape[0])
        for i in range(v1.shape[0]):
            ans[i] = func(v1[i],v2[i])
        return ans

def _findMax(a,b):
    if a>b:
        return a
    else:
        return b

def _findMin(a,b):
    if a<b:
        return a
    else:
        return b


# if __name__ == '__main__':
#     a = np.array([
#         [0,0,1,1],
#         [1,1,2,2],
#         [3,3,4,4]])
#     b = np.array([
#         [0.5,0.5,1.5,1.5],
#         [0.5,0.5,1.5,1.5],
#         [0,0,1,1]])
#     print(getIouArray(a,b))

#     a1 = (0,0,1,1)
#     a2 = (1,1,2,2)
#     a3 = (3,3,4,4)

#     b1 = (0.5,0.5,1.5,1.5)
#     b2 = (0.5,0.5,1.5,1.5)
#     b3 = (0,0,1,1)
#     print(get_iou(a1,b1))
#     print(get_iou(a2,b2))
#     print(get_iou(a3,b3))